Gaussian Processes

Regression
Non-parametric
A Bayesian approach to regression and classification that defines a distribution over functions.

General Principles

Through varying intercepts and slopes, we have seen how to quantify some of the unique features that generate variation across clusters and covariance among the observations within each cluster. But through the covariance matrix that is used to account for correlation between clusters, we are inherently assuming linear relationships between clusters. What if we want to model the relationship between two variables that are not linearly related? In this case, we can use a Gaussian Process (GP) to model the relationship between two variables.

Considerations

Caution
  • To capture complex, non-linear relationships in data where the underlying function is smooth but has an unknown functional form, GPs use a kernel πŸ›ˆ.
  • The choice of kernel hyperparameters can significantly impact results; thus, GPs require choosing an appropriate kernel function that captures the expected behavior of your data.
  • Through kernel definition, we can incorporate domain knowledge.
  • They scale poorly with dataset size (O(nΒ³) complexity) due to matrix operations; thus, memory requirements can be substantial for large datasets, which has led to neural networks being used instead to resolve large non-linear problems.

Example

Below is an example code snippet demonstrating Gaussian Process regression using the Bayesian Inference (BI) package. Data consist of a continuous dependent variable (total_tools), representing the number of tools invented in the islands, and a continuous independent variable (population), representing the population of the islands. The goal is to estimate the effect of population on the total tools. We use the distance matrix of the islands for the kernel function in order to capture the spatial dependence of the relationship. This example is based on McElreath (2018).

Code 
from BI import bi, jnp
import pandas as pd
# Setup device------------------------------------------------
m = bi(platform='cpu')

# Import Data & Data Manipulation ------------------------------------------------
# Import
from importlib.resources import files
data_path = m.load.kline2(only_path=True)
m.data(data_path, sep=';') 


data_path2 = files('BI.Resources') / 'islandsDistMatrix.csv'
islandsDistMatrix = pd.read_csv(data_path2, index_col=0)

m.data_to_model(['total_tools', 'population'])
m.data_on_model["society"] = jnp.arange(0,10)# index observations
m.data_on_model["Dmat"] = islandsDistMatrix.values # Distance matrix

def model(Dmat, population, society, total_tools):
    a = m.dist.exponential(1, name = 'a')
    b = m.dist.exponential(1, name = 'b')
    g = m.dist.exponential(1, name = 'g')

    # non-centered Gaussian Process prior
    etasq = m.dist.exponential(2, name = 'etasq')
    rhosq = m.dist.exponential(0.5, name = 'rhosq')
    SIGMA = etasq * jnp.exp(-rhosq * jnp.square(Dmat))
    SIGMA = SIGMA.at[jnp.diag_indices(Dmat.shape[0])].add(etasq)
    k = m.dist.multivariate_normal(0, SIGMA, name = 'k')

    lambda_ = a * population**b / g * jnp.exp(k[society])

    m.dist.poisson(lambda_, obs=total_tools)

# Run sampler ------------------------------------------------
m.fit(model) 
m.summary()
jax.local_device_count 16
  0%|          | 0/1000 [00:00<?, ?it/s]warmup:   0%|          | 1/1000 [00:04<1:15:12,  4.52s/it, 1 steps of size 2.34e+00. acc. prob=0.00]warmup:   3%|β–Ž         | 26/1000 [00:04<02:04,  7.85it/s, 255 steps of size 1.44e-02. acc. prob=0.71]warmup:   4%|▍         | 39/1000 [00:04<01:16, 12.59it/s, 127 steps of size 1.80e-02. acc. prob=0.74]warmup:   5%|β–Œ         | 50/1000 [00:05<00:59, 16.09it/s, 95 steps of size 3.19e-02. acc. prob=0.75] warmup:   6%|β–Œ         | 58/1000 [00:05<00:47, 19.93it/s, 191 steps of size 1.58e-02. acc. prob=0.75]warmup:   7%|β–‹         | 66/1000 [00:05<00:37, 24.85it/s, 127 steps of size 5.67e-03. acc. prob=0.75]warmup:   7%|β–‹         | 74/1000 [00:05<00:33, 28.06it/s, 95 steps of size 2.62e-02. acc. prob=0.76] warmup:   8%|β–Š         | 81/1000 [00:05<00:29, 31.13it/s, 127 steps of size 2.29e-02. acc. prob=0.76]warmup:   9%|β–Š         | 87/1000 [00:05<00:26, 34.76it/s, 127 steps of size 2.42e-02. acc. prob=0.76]warmup:  10%|β–‰         | 99/1000 [00:05<00:18, 48.75it/s, 95 steps of size 2.46e-02. acc. prob=0.77] warmup:  11%|β–ˆ         | 107/1000 [00:06<00:17, 50.52it/s, 63 steps of size 7.46e-02. acc. prob=0.76]warmup:  11%|β–ˆβ–        | 114/1000 [00:06<00:19, 44.90it/s, 383 steps of size 2.25e-02. acc. prob=0.76]warmup:  12%|β–ˆβ–        | 120/1000 [00:06<00:18, 46.66it/s, 255 steps of size 1.67e-02. acc. prob=0.76]warmup:  13%|β–ˆβ–Ž        | 126/1000 [00:06<00:19, 44.81it/s, 5 steps of size 8.79e-03. acc. prob=0.76]  warmup:  13%|β–ˆβ–Ž        | 132/1000 [00:06<00:23, 37.47it/s, 127 steps of size 2.46e-02. acc. prob=0.77]warmup:  14%|β–ˆβ–        | 138/1000 [00:06<00:20, 41.10it/s, 127 steps of size 3.43e-02. acc. prob=0.77]warmup:  14%|β–ˆβ–        | 145/1000 [00:06<00:19, 44.54it/s, 255 steps of size 2.95e-02. acc. prob=0.77]warmup:  15%|β–ˆβ–Œ        | 151/1000 [00:07<00:17, 47.68it/s, 31 steps of size 2.53e-01. acc. prob=0.77] warmup:  16%|β–ˆβ–Œ        | 157/1000 [00:07<00:18, 45.93it/s, 127 steps of size 6.04e-02. acc. prob=0.77]warmup:  16%|β–ˆβ–Œ        | 162/1000 [00:07<00:19, 43.02it/s, 127 steps of size 4.11e-02. acc. prob=0.77]warmup:  17%|β–ˆβ–‹        | 167/1000 [00:07<00:25, 32.74it/s, 255 steps of size 4.43e-02. acc. prob=0.77]warmup:  17%|β–ˆβ–‹        | 171/1000 [00:07<00:26, 31.79it/s, 511 steps of size 1.79e-02. acc. prob=0.77]warmup:  18%|β–ˆβ–Š        | 175/1000 [00:07<00:24, 33.43it/s, 6 steps of size 9.64e-03. acc. prob=0.77]  warmup:  18%|β–ˆβ–Š        | 179/1000 [00:08<00:31, 25.79it/s, 63 steps of size 3.18e-02. acc. prob=0.77]warmup:  19%|β–ˆβ–Š        | 186/1000 [00:08<00:23, 33.95it/s, 255 steps of size 3.61e-02. acc. prob=0.77]warmup:  19%|β–ˆβ–‰        | 191/1000 [00:08<00:23, 34.15it/s, 255 steps of size 4.68e-02. acc. prob=0.77]warmup:  20%|β–ˆβ–‰        | 195/1000 [00:08<00:26, 30.30it/s, 255 steps of size 5.64e-02. acc. prob=0.77]warmup:  20%|β–ˆβ–ˆ        | 200/1000 [00:08<00:23, 33.81it/s, 127 steps of size 3.84e-02. acc. prob=0.77]warmup:  20%|β–ˆβ–ˆ        | 205/1000 [00:08<00:22, 34.98it/s, 255 steps of size 3.04e-02. acc. prob=0.77]warmup:  21%|β–ˆβ–ˆ        | 209/1000 [00:08<00:22, 34.47it/s, 127 steps of size 5.53e-02. acc. prob=0.78]warmup:  21%|β–ˆβ–ˆβ–       | 213/1000 [00:09<00:22, 35.04it/s, 127 steps of size 4.42e-02. acc. prob=0.78]warmup:  22%|β–ˆβ–ˆβ–       | 218/1000 [00:09<00:21, 35.66it/s, 255 steps of size 4.07e-02. acc. prob=0.78]warmup:  22%|β–ˆβ–ˆβ–       | 222/1000 [00:09<00:24, 32.30it/s, 511 steps of size 3.12e-02. acc. prob=0.78]warmup:  23%|β–ˆβ–ˆβ–Ž       | 227/1000 [00:09<00:21, 36.29it/s, 127 steps of size 4.57e-02. acc. prob=0.78]warmup:  23%|β–ˆβ–ˆβ–Ž       | 232/1000 [00:09<00:20, 37.88it/s, 127 steps of size 6.04e-02. acc. prob=0.78]warmup:  24%|β–ˆβ–ˆβ–Ž       | 237/1000 [00:09<00:19, 38.41it/s, 127 steps of size 6.16e-02. acc. prob=0.78]warmup:  24%|β–ˆβ–ˆβ–       | 241/1000 [00:09<00:22, 33.26it/s, 255 steps of size 4.26e-02. acc. prob=0.78]warmup:  24%|β–ˆβ–ˆβ–       | 245/1000 [00:09<00:22, 33.08it/s, 63 steps of size 5.84e-02. acc. prob=0.78] warmup:  25%|β–ˆβ–ˆβ–       | 249/1000 [00:10<00:23, 32.32it/s, 127 steps of size 5.20e-02. acc. prob=0.78]warmup:  25%|β–ˆβ–ˆβ–Œ       | 253/1000 [00:10<00:23, 31.97it/s, 127 steps of size 1.56e-02. acc. prob=0.77]warmup:  26%|β–ˆβ–ˆβ–Œ       | 257/1000 [00:10<00:25, 29.60it/s, 63 steps of size 7.81e-02. acc. prob=0.78] warmup:  26%|β–ˆβ–ˆβ–Œ       | 261/1000 [00:10<00:25, 29.47it/s, 127 steps of size 4.11e-02. acc. prob=0.78]warmup:  26%|β–ˆβ–ˆβ–‹       | 264/1000 [00:10<00:28, 26.01it/s, 511 steps of size 1.65e-02. acc. prob=0.78]warmup:  27%|β–ˆβ–ˆβ–‹       | 267/1000 [00:10<00:28, 25.75it/s, 31 steps of size 3.86e-02. acc. prob=0.78] warmup:  27%|β–ˆβ–ˆβ–‹       | 270/1000 [00:10<00:27, 26.30it/s, 255 steps of size 3.00e-02. acc. prob=0.78]warmup:  27%|β–ˆβ–ˆβ–‹       | 274/1000 [00:11<00:26, 27.79it/s, 127 steps of size 3.73e-02. acc. prob=0.78]warmup:  28%|β–ˆβ–ˆβ–Š       | 278/1000 [00:11<00:25, 28.83it/s, 255 steps of size 1.97e-02. acc. prob=0.78]warmup:  28%|β–ˆβ–ˆβ–Š       | 282/1000 [00:11<00:24, 29.67it/s, 127 steps of size 1.50e-02. acc. prob=0.78]warmup:  28%|β–ˆβ–ˆβ–Š       | 285/1000 [00:11<00:24, 29.15it/s, 127 steps of size 5.82e-02. acc. prob=0.78]warmup:  29%|β–ˆβ–ˆβ–‰       | 289/1000 [00:11<00:23, 30.28it/s, 127 steps of size 3.76e-02. acc. prob=0.78]warmup:  29%|β–ˆβ–ˆβ–‰       | 294/1000 [00:11<00:20, 34.78it/s, 127 steps of size 4.89e-02. acc. prob=0.78]warmup:  30%|β–ˆβ–ˆβ–ˆ       | 300/1000 [00:11<00:17, 39.28it/s, 127 steps of size 3.27e-02. acc. prob=0.78]warmup:  30%|β–ˆβ–ˆβ–ˆ       | 305/1000 [00:11<00:17, 40.37it/s, 127 steps of size 3.46e-02. acc. prob=0.78]warmup:  31%|β–ˆβ–ˆβ–ˆ       | 311/1000 [00:11<00:15, 44.26it/s, 127 steps of size 3.31e-02. acc. prob=0.78]warmup:  32%|β–ˆβ–ˆβ–ˆβ–      | 316/1000 [00:12<00:15, 45.25it/s, 127 steps of size 3.80e-02. acc. prob=0.78]warmup:  32%|β–ˆβ–ˆβ–ˆβ–      | 321/1000 [00:12<00:15, 43.86it/s, 63 steps of size 4.73e-02. acc. prob=0.78] warmup:  33%|β–ˆβ–ˆβ–ˆβ–Ž      | 328/1000 [00:12<00:13, 50.79it/s, 63 steps of size 6.13e-02. acc. prob=0.78]warmup:  33%|β–ˆβ–ˆβ–ˆβ–Ž      | 334/1000 [00:12<00:13, 48.53it/s, 255 steps of size 3.03e-02. acc. prob=0.78]warmup:  34%|β–ˆβ–ˆβ–ˆβ–      | 339/1000 [00:12<00:13, 48.25it/s, 127 steps of size 2.61e-02. acc. prob=0.78]warmup:  34%|β–ˆβ–ˆβ–ˆβ–      | 345/1000 [00:12<00:13, 49.39it/s, 383 steps of size 2.52e-02. acc. prob=0.78]warmup:  35%|β–ˆβ–ˆβ–ˆβ–Œ      | 351/1000 [00:12<00:12, 50.91it/s, 63 steps of size 7.18e-02. acc. prob=0.78] warmup:  36%|β–ˆβ–ˆβ–ˆβ–Œ      | 357/1000 [00:12<00:13, 46.81it/s, 127 steps of size 3.88e-02. acc. prob=0.78]warmup:  37%|β–ˆβ–ˆβ–ˆβ–‹      | 367/1000 [00:12<00:10, 58.96it/s, 127 steps of size 4.95e-02. acc. prob=0.78]warmup:  38%|β–ˆβ–ˆβ–ˆβ–Š      | 378/1000 [00:13<00:08, 71.81it/s, 63 steps of size 5.89e-02. acc. prob=0.78] warmup:  39%|β–ˆβ–ˆβ–ˆβ–‰      | 388/1000 [00:13<00:07, 78.71it/s, 63 steps of size 4.71e-02. acc. prob=0.78]warmup:  40%|β–ˆβ–ˆβ–ˆβ–ˆ      | 404/1000 [00:13<00:05, 100.92it/s, 127 steps of size 5.14e-02. acc. prob=0.78]warmup:  42%|β–ˆβ–ˆβ–ˆβ–ˆβ–     | 415/1000 [00:13<00:06, 91.98it/s, 63 steps of size 4.13e-02. acc. prob=0.78]  warmup:  42%|β–ˆβ–ˆβ–ˆβ–ˆβ–Ž     | 425/1000 [00:13<00:07, 72.91it/s, 63 steps of size 5.87e-02. acc. prob=0.78]warmup:  43%|β–ˆβ–ˆβ–ˆβ–ˆβ–Ž     | 434/1000 [00:13<00:09, 57.40it/s, 127 steps of size 5.95e-02. acc. prob=0.78]warmup:  44%|β–ˆβ–ˆβ–ˆβ–ˆβ–     | 441/1000 [00:14<00:14, 38.45it/s, 63 steps of size 7.60e-02. acc. prob=0.78] warmup:  45%|β–ˆβ–ˆβ–ˆβ–ˆβ–     | 447/1000 [00:14<00:13, 39.69it/s, 127 steps of size 3.87e-02. acc. prob=0.78]warmup:  46%|β–ˆβ–ˆβ–ˆβ–ˆβ–Œ     | 455/1000 [00:14<00:12, 42.13it/s, 255 steps of size 2.98e-02. acc. prob=0.78]warmup:  46%|β–ˆβ–ˆβ–ˆβ–ˆβ–Œ     | 461/1000 [00:15<00:20, 26.88it/s, 255 steps of size 3.79e-02. acc. prob=0.78]warmup:  46%|β–ˆβ–ˆβ–ˆβ–ˆβ–‹     | 465/1000 [00:15<00:18, 28.31it/s, 127 steps of size 3.49e-02. acc. prob=0.78]warmup:  47%|β–ˆβ–ˆβ–ˆβ–ˆβ–‹     | 469/1000 [00:15<00:25, 20.68it/s, 511 steps of size 2.14e-02. acc. prob=0.78]warmup:  47%|β–ˆβ–ˆβ–ˆβ–ˆβ–‹     | 472/1000 [00:15<00:24, 21.59it/s, 35 steps of size 2.07e-02. acc. prob=0.78] warmup:  48%|β–ˆβ–ˆβ–ˆβ–ˆβ–Š     | 476/1000 [00:15<00:22, 23.81it/s, 63 steps of size 8.16e-02. acc. prob=0.78]warmup:  48%|β–ˆβ–ˆβ–ˆβ–ˆβ–Š     | 480/1000 [00:15<00:22, 23.25it/s, 63 steps of size 4.78e-02. acc. prob=0.78]warmup:  48%|β–ˆβ–ˆβ–ˆβ–ˆβ–Š     | 485/1000 [00:16<00:22, 23.30it/s, 511 steps of size 1.52e-02. acc. prob=0.78]warmup:  49%|β–ˆβ–ˆβ–ˆβ–ˆβ–‰     | 488/1000 [00:16<00:22, 22.92it/s, 63 steps of size 7.24e-02. acc. prob=0.78] warmup:  49%|β–ˆβ–ˆβ–ˆβ–ˆβ–‰     | 491/1000 [00:16<00:22, 22.19it/s, 255 steps of size 2.88e-02. acc. prob=0.78]warmup:  49%|β–ˆβ–ˆβ–ˆβ–ˆβ–‰     | 494/1000 [00:16<00:21, 23.13it/s, 127 steps of size 4.66e-02. acc. prob=0.78]warmup:  50%|β–ˆβ–ˆβ–ˆβ–ˆβ–‰     | 497/1000 [00:16<00:21, 23.67it/s, 255 steps of size 3.26e-02. acc. prob=0.78]sample:  50%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆ     | 502/1000 [00:16<00:18, 26.54it/s, 127 steps of size 3.06e-02. acc. prob=0.97]sample:  51%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆ     | 506/1000 [00:16<00:17, 27.94it/s, 127 steps of size 3.06e-02. acc. prob=0.97]sample:  51%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆ     | 509/1000 [00:17<00:17, 27.33it/s, 127 steps of size 3.06e-02. acc. prob=0.97]sample:  51%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆ     | 512/1000 [00:17<00:17, 27.88it/s, 127 steps of size 3.06e-02. acc. prob=0.97]sample:  52%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–    | 515/1000 [00:17<00:17, 27.24it/s, 127 steps of size 3.06e-02. acc. prob=0.96]sample:  52%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–    | 518/1000 [00:17<00:18, 26.15it/s, 127 steps of size 3.06e-02. acc. prob=0.97]sample:  52%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–    | 521/1000 [00:17<00:18, 25.88it/s, 63 steps of size 3.06e-02. acc. prob=0.94] sample:  52%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–    | 524/1000 [00:17<00:21, 22.27it/s, 255 steps of size 3.06e-02. acc. prob=0.95]sample:  53%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž    | 527/1000 [00:17<00:21, 22.09it/s, 63 steps of size 3.06e-02. acc. prob=0.94] sample:  53%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž    | 530/1000 [00:18<00:21, 21.58it/s, 127 steps of size 3.06e-02. acc. prob=0.95]sample:  53%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž    | 533/1000 [00:18<00:20, 22.78it/s, 127 steps of size 3.06e-02. acc. prob=0.95]sample:  54%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž    | 537/1000 [00:18<00:18, 25.02it/s, 127 steps of size 3.06e-02. acc. prob=0.95]sample:  54%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–    | 541/1000 [00:18<00:16, 27.84it/s, 63 steps of size 3.06e-02. acc. prob=0.94] sample:  55%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–    | 545/1000 [00:18<00:15, 29.77it/s, 127 steps of size 3.06e-02. acc. prob=0.93]sample:  55%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–    | 549/1000 [00:18<00:15, 29.70it/s, 127 steps of size 3.06e-02. acc. prob=0.93]sample:  55%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ    | 553/1000 [00:18<00:15, 28.93it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  56%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ    | 557/1000 [00:18<00:14, 29.95it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  56%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ    | 561/1000 [00:19<00:15, 27.74it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  56%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹    | 564/1000 [00:19<00:15, 27.34it/s, 127 steps of size 3.06e-02. acc. prob=0.93]sample:  57%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹    | 568/1000 [00:19<00:15, 28.55it/s, 127 steps of size 3.06e-02. acc. prob=0.93]sample:  57%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹    | 572/1000 [00:19<00:15, 27.49it/s, 255 steps of size 3.06e-02. acc. prob=0.93]sample:  57%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š    | 575/1000 [00:19<00:15, 27.75it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  58%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š    | 578/1000 [00:19<00:15, 27.49it/s, 127 steps of size 3.06e-02. acc. prob=0.93]sample:  58%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š    | 582/1000 [00:19<00:14, 29.65it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  59%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š    | 587/1000 [00:19<00:13, 31.72it/s, 127 steps of size 3.06e-02. acc. prob=0.93]sample:  59%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰    | 591/1000 [00:20<00:15, 26.04it/s, 255 steps of size 3.06e-02. acc. prob=0.94]sample:  60%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰    | 595/1000 [00:20<00:15, 26.88it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  60%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰    | 598/1000 [00:20<00:16, 24.22it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  60%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ    | 601/1000 [00:20<00:17, 22.42it/s, 127 steps of size 3.06e-02. acc. prob=0.93]sample:  60%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ    | 604/1000 [00:20<00:17, 22.59it/s, 255 steps of size 3.06e-02. acc. prob=0.93]sample:  61%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ    | 608/1000 [00:20<00:15, 25.25it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  61%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ    | 611/1000 [00:21<00:16, 23.87it/s, 255 steps of size 3.06e-02. acc. prob=0.94]sample:  61%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–   | 614/1000 [00:21<00:15, 24.82it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  62%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–   | 617/1000 [00:21<00:16, 23.39it/s, 191 steps of size 3.06e-02. acc. prob=0.94]sample:  62%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–   | 622/1000 [00:21<00:13, 28.05it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  63%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž   | 626/1000 [00:21<00:13, 26.98it/s, 255 steps of size 3.06e-02. acc. prob=0.94]sample:  63%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž   | 629/1000 [00:21<00:14, 25.82it/s, 255 steps of size 3.06e-02. acc. prob=0.94]sample:  63%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž   | 633/1000 [00:21<00:13, 27.20it/s, 255 steps of size 3.06e-02. acc. prob=0.94]sample:  64%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž   | 637/1000 [00:21<00:12, 28.46it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  64%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–   | 641/1000 [00:22<00:12, 28.89it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  64%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–   | 645/1000 [00:22<00:11, 30.40it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  65%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–   | 649/1000 [00:22<00:11, 30.12it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  65%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ   | 653/1000 [00:22<00:11, 31.30it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  66%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ   | 657/1000 [00:22<00:10, 31.91it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  66%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ   | 661/1000 [00:22<00:11, 29.92it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  66%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹   | 665/1000 [00:22<00:11, 30.42it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  67%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹   | 669/1000 [00:22<00:10, 30.73it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  67%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹   | 673/1000 [00:23<00:11, 28.82it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  68%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š   | 676/1000 [00:23<00:11, 28.93it/s, 63 steps of size 3.06e-02. acc. prob=0.94] sample:  68%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š   | 679/1000 [00:23<00:12, 24.88it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  68%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š   | 682/1000 [00:23<00:12, 25.83it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  68%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š   | 685/1000 [00:23<00:12, 25.26it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  69%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰   | 688/1000 [00:23<00:12, 24.77it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  69%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰   | 691/1000 [00:23<00:11, 25.93it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  69%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰   | 694/1000 [00:24<00:13, 23.22it/s, 191 steps of size 3.06e-02. acc. prob=0.94]sample:  70%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰   | 697/1000 [00:24<00:12, 24.35it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  70%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ   | 700/1000 [00:24<00:11, 25.42it/s, 63 steps of size 3.06e-02. acc. prob=0.94] sample:  70%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ   | 703/1000 [00:24<00:11, 25.20it/s, 63 steps of size 3.06e-02. acc. prob=0.94]sample:  71%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ   | 706/1000 [00:24<00:11, 25.36it/s, 63 steps of size 3.06e-02. acc. prob=0.94]sample:  71%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ   | 709/1000 [00:24<00:11, 24.49it/s, 255 steps of size 3.06e-02. acc. prob=0.94]sample:  71%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–  | 713/1000 [00:24<00:10, 26.54it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  72%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–  | 716/1000 [00:24<00:10, 26.51it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  72%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–  | 721/1000 [00:24<00:09, 30.89it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  72%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž  | 725/1000 [00:25<00:09, 28.38it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  73%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž  | 728/1000 [00:25<00:10, 26.99it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  73%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž  | 732/1000 [00:25<00:09, 28.74it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  74%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž  | 737/1000 [00:25<00:07, 32.88it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  74%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–  | 742/1000 [00:25<00:06, 37.04it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  75%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–  | 748/1000 [00:25<00:05, 42.60it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  75%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ  | 754/1000 [00:25<00:05, 45.00it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  76%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ  | 761/1000 [00:25<00:04, 51.28it/s, 63 steps of size 3.06e-02. acc. prob=0.94] sample:  77%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹  | 768/1000 [00:26<00:04, 55.27it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  77%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹  | 774/1000 [00:26<00:04, 50.84it/s, 127 steps of size 3.06e-02. acc. prob=0.95]sample:  78%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š  | 780/1000 [00:26<00:04, 49.00it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  79%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š  | 787/1000 [00:26<00:04, 53.16it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  79%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰  | 794/1000 [00:26<00:03, 57.35it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  80%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ  | 802/1000 [00:26<00:03, 61.81it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  81%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ  | 809/1000 [00:26<00:03, 56.14it/s, 63 steps of size 3.06e-02. acc. prob=0.94] sample:  82%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 815/1000 [00:26<00:03, 54.16it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  82%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 821/1000 [00:27<00:03, 54.66it/s, 63 steps of size 3.06e-02. acc. prob=0.94] sample:  83%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 829/1000 [00:27<00:02, 58.52it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  84%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 836/1000 [00:27<00:02, 59.88it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  84%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 843/1000 [00:27<00:02, 54.76it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  85%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 850/1000 [00:27<00:02, 58.03it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  86%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 862/1000 [00:27<00:01, 72.96it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  87%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹ | 873/1000 [00:27<00:01, 81.17it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  88%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š | 883/1000 [00:27<00:01, 85.99it/s, 255 steps of size 3.06e-02. acc. prob=0.94]sample:  89%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 894/1000 [00:27<00:01, 92.58it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  90%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 904/1000 [00:28<00:01, 93.05it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  91%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–| 914/1000 [00:28<00:00, 94.00it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  92%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–| 924/1000 [00:28<00:00, 94.57it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  93%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž| 934/1000 [00:28<00:00, 94.82it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  94%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–| 945/1000 [00:28<00:00, 99.14it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  96%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ| 956/1000 [00:28<00:00, 99.45it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  97%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹| 967/1000 [00:28<00:00, 99.62it/s, 255 steps of size 3.06e-02. acc. prob=0.94]sample:  98%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š| 977/1000 [00:28<00:00, 83.20it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample:  99%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰| 988/1000 [00:28<00:00, 88.88it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰| 998/1000 [00:29<00:00, 90.90it/s, 127 steps of size 3.06e-02. acc. prob=0.94]sample: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 1000/1000 [00:29<00:00, 34.40it/s, 127 steps of size 3.06e-02. acc. prob=0.94]
arviz - WARNING - Shape validation failed: input_shape: (1, 500), minimum_shape: (chains=2, draws=4)
mean sd hdi_5.5% hdi_94.5% mcse_mean mcse_sd ess_bulk ess_tail r_hat
a 1.40 1.03 0.20 2.82 0.07 0.07 162.10 261.31 NaN
b 0.29 0.08 0.17 0.41 0.01 0.00 172.66 189.40 NaN
etasq 0.08 0.05 0.01 0.15 0.00 0.00 148.10 193.35 NaN
g 0.65 0.63 0.01 1.35 0.04 0.07 190.39 243.19 NaN
k[0] -0.20 0.28 -0.60 0.26 0.02 0.02 176.77 195.46 NaN
k[1] 0.03 0.26 -0.37 0.42 0.02 0.01 214.62 210.07 NaN
k[2] -0.06 0.24 -0.43 0.30 0.02 0.01 182.16 196.83 NaN
k[3] 0.36 0.21 0.05 0.71 0.02 0.01 165.96 179.87 NaN
k[4] 0.05 0.21 -0.32 0.35 0.02 0.01 152.21 146.88 NaN
k[5] -0.37 0.25 -0.77 -0.01 0.02 0.01 153.15 217.39 NaN
k[6] 0.14 0.19 -0.20 0.43 0.02 0.01 123.67 208.08 NaN
k[7] -0.21 0.22 -0.55 0.15 0.02 0.01 145.49 288.18 NaN
k[8] 0.26 0.20 -0.03 0.59 0.02 0.01 141.54 171.81 NaN
k[9] -0.20 0.30 -0.68 0.24 0.03 0.02 131.02 243.52 NaN
rhosq 1.62 1.82 0.00 3.20 0.10 0.17 199.68 165.94 NaN
Code 
from BI import bi, jnp
import pandas as pd
# Setup device------------------------------------------------
m = bi(platform='cpu')

# Import Data & Data Manipulation ------------------------------------------------
# Import
from importlib.resources import files
data_path = m.load.kline2(only_path=True)
m.data(data_path, sep=';') 

islandsDistMatrix = m.load.islands_dist_matrix(frame = False)['data']

m.data_to_model(['total_tools', 'population'])
m.data_on_model["society"] = jnp.arange(0,10)# index observations
m.data_on_model["Dmat"] = islandsDistMatrix # Distance matrix


def model(Dmat, population, society, total_tools):
    a = m.dist.exponential(1, name = 'a')
    b = m.dist.exponential(1, name = 'b')
    g = m.dist.exponential(1, name = 'g')

    k = m.gaussian.gaussian_process(Dmat)

    lambda_ = a * population**b / g * jnp.exp(k[society])

    m.dist.poisson(lambda_, obs=total_tools)

# Run sampler ------------------------------------------------
m.fit(model) 
m.summary()
jax.local_device_count 16
  0%|          | 0/1000 [00:00<?, ?it/s]warmup:   0%|          | 1/1000 [00:04<1:14:31,  4.48s/it, 1 steps of size 2.34e+00. acc. prob=0.00]warmup:   3%|β–Ž         | 31/1000 [00:04<01:42,  9.48it/s, 31 steps of size 1.29e-02. acc. prob=0.72]warmup:   5%|▍         | 49/1000 [00:04<00:57, 16.57it/s, 159 steps of size 1.74e-02. acc. prob=0.75]warmup:   7%|β–‹         | 66/1000 [00:04<00:37, 24.71it/s, 127 steps of size 2.99e-02. acc. prob=0.76]warmup:   8%|β–Š         | 81/1000 [00:04<00:27, 33.05it/s, 127 steps of size 1.77e-02. acc. prob=0.76]warmup:  10%|β–‰         | 95/1000 [00:05<00:22, 41.07it/s, 255 steps of size 2.07e-02. acc. prob=0.76]warmup:  11%|β–ˆ         | 108/1000 [00:05<00:17, 50.42it/s, 191 steps of size 2.88e-02. acc. prob=0.76]warmup:  12%|β–ˆβ–        | 122/1000 [00:05<00:14, 62.39it/s, 63 steps of size 4.92e-02. acc. prob=0.77] warmup:  14%|β–ˆβ–Ž        | 135/1000 [00:05<00:12, 71.02it/s, 191 steps of size 3.10e-02. acc. prob=0.77]warmup:  15%|β–ˆβ–        | 147/1000 [00:05<00:10, 78.14it/s, 127 steps of size 7.39e-02. acc. prob=0.77]warmup:  16%|β–ˆβ–Œ        | 159/1000 [00:05<00:10, 81.43it/s, 127 steps of size 3.46e-02. acc. prob=0.77]warmup:  17%|β–ˆβ–‹        | 170/1000 [00:05<00:10, 75.97it/s, 57 steps of size 5.88e-02. acc. prob=0.77] warmup:  18%|β–ˆβ–Š        | 180/1000 [00:06<00:11, 73.64it/s, 511 steps of size 2.03e-02. acc. prob=0.77]warmup:  19%|β–ˆβ–‰        | 193/1000 [00:06<00:09, 84.54it/s, 127 steps of size 3.35e-02. acc. prob=0.77]warmup:  20%|β–ˆβ–ˆ        | 203/1000 [00:06<00:09, 83.29it/s, 127 steps of size 6.34e-02. acc. prob=0.78]warmup:  22%|β–ˆβ–ˆβ–       | 219/1000 [00:06<00:07, 98.09it/s, 255 steps of size 3.68e-02. acc. prob=0.78]warmup:  24%|β–ˆβ–ˆβ–Ž       | 235/1000 [00:06<00:06, 112.63it/s, 127 steps of size 3.61e-02. acc. prob=0.78]warmup:  25%|β–ˆβ–ˆβ–       | 248/1000 [00:06<00:06, 110.60it/s, 63 steps of size 5.24e-02. acc. prob=0.78] warmup:  26%|β–ˆβ–ˆβ–‹       | 265/1000 [00:06<00:05, 122.92it/s, 255 steps of size 1.71e-02. acc. prob=0.78]warmup:  28%|β–ˆβ–ˆβ–Š       | 281/1000 [00:06<00:05, 131.89it/s, 127 steps of size 5.27e-02. acc. prob=0.78]warmup:  30%|β–ˆβ–ˆβ–‰       | 297/1000 [00:06<00:05, 138.03it/s, 63 steps of size 5.29e-02. acc. prob=0.78] warmup:  31%|β–ˆβ–ˆβ–ˆ       | 312/1000 [00:07<00:05, 137.39it/s, 511 steps of size 3.37e-02. acc. prob=0.78]warmup:  33%|β–ˆβ–ˆβ–ˆβ–Ž      | 326/1000 [00:07<00:05, 126.56it/s, 63 steps of size 4.41e-02. acc. prob=0.78] warmup:  34%|β–ˆβ–ˆβ–ˆβ–      | 342/1000 [00:07<00:05, 131.31it/s, 191 steps of size 5.09e-02. acc. prob=0.78]warmup:  36%|β–ˆβ–ˆβ–ˆβ–Œ      | 356/1000 [00:07<00:04, 130.66it/s, 127 steps of size 3.18e-02. acc. prob=0.78]warmup:  37%|β–ˆβ–ˆβ–ˆβ–‹      | 372/1000 [00:07<00:04, 136.52it/s, 63 steps of size 3.87e-02. acc. prob=0.78] warmup:  39%|β–ˆβ–ˆβ–ˆβ–Š      | 386/1000 [00:07<00:04, 135.21it/s, 43 steps of size 3.02e-02. acc. prob=0.78]warmup:  40%|β–ˆβ–ˆβ–ˆβ–ˆ      | 400/1000 [00:07<00:04, 126.66it/s, 127 steps of size 3.70e-02. acc. prob=0.78]warmup:  41%|β–ˆβ–ˆβ–ˆβ–ˆβ–     | 413/1000 [00:07<00:04, 122.50it/s, 5 steps of size 2.61e-02. acc. prob=0.78]  warmup:  43%|β–ˆβ–ˆβ–ˆβ–ˆβ–Ž     | 433/1000 [00:07<00:04, 141.72it/s, 159 steps of size 4.96e-02. acc. prob=0.78]warmup:  45%|β–ˆβ–ˆβ–ˆβ–ˆβ–     | 448/1000 [00:08<00:03, 138.50it/s, 127 steps of size 5.64e-02. acc. prob=0.78]warmup:  46%|β–ˆβ–ˆβ–ˆβ–ˆβ–‹     | 463/1000 [00:08<00:04, 123.22it/s, 31 steps of size 1.13e-01. acc. prob=0.78] warmup:  48%|β–ˆβ–ˆβ–ˆβ–ˆβ–Š     | 476/1000 [00:08<00:04, 112.40it/s, 63 steps of size 8.72e-02. acc. prob=0.78]warmup:  49%|β–ˆβ–ˆβ–ˆβ–ˆβ–‰     | 488/1000 [00:08<00:04, 105.39it/s, 5 steps of size 1.46e-02. acc. prob=0.78] sample:  50%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆ     | 503/1000 [00:08<00:04, 115.61it/s, 127 steps of size 3.34e-02. acc. prob=0.99]sample:  52%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–    | 515/1000 [00:08<00:04, 112.82it/s, 63 steps of size 3.34e-02. acc. prob=0.95] sample:  53%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž    | 527/1000 [00:08<00:04, 108.65it/s, 63 steps of size 3.34e-02. acc. prob=0.94]sample:  54%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–    | 539/1000 [00:08<00:04, 109.39it/s, 31 steps of size 3.34e-02. acc. prob=0.93]sample:  56%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ    | 556/1000 [00:09<00:03, 123.86it/s, 191 steps of size 3.34e-02. acc. prob=0.93]sample:  57%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹    | 572/1000 [00:09<00:03, 133.42it/s, 63 steps of size 3.34e-02. acc. prob=0.93] sample:  59%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰    | 589/1000 [00:09<00:02, 142.28it/s, 63 steps of size 3.34e-02. acc. prob=0.93]sample:  60%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ    | 604/1000 [00:09<00:02, 132.11it/s, 127 steps of size 3.34e-02. acc. prob=0.93]sample:  62%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–   | 618/1000 [00:09<00:02, 131.19it/s, 127 steps of size 3.34e-02. acc. prob=0.93]sample:  63%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž   | 634/1000 [00:09<00:02, 135.34it/s, 127 steps of size 3.34e-02. acc. prob=0.93]sample:  65%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–   | 648/1000 [00:09<00:02, 130.06it/s, 63 steps of size 3.34e-02. acc. prob=0.94] sample:  66%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹   | 663/1000 [00:09<00:02, 135.44it/s, 63 steps of size 3.34e-02. acc. prob=0.94]sample:  68%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š   | 677/1000 [00:09<00:02, 128.93it/s, 63 steps of size 3.34e-02. acc. prob=0.94]sample:  69%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰   | 691/1000 [00:10<00:02, 125.67it/s, 255 steps of size 3.34e-02. acc. prob=0.94]sample:  70%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ   | 704/1000 [00:10<00:02, 119.80it/s, 255 steps of size 3.34e-02. acc. prob=0.94]sample:  72%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–  | 717/1000 [00:10<00:02, 116.68it/s, 63 steps of size 3.34e-02. acc. prob=0.93] sample:  73%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž  | 730/1000 [00:10<00:02, 119.29it/s, 63 steps of size 3.34e-02. acc. prob=0.93]sample:  74%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–  | 745/1000 [00:10<00:02, 125.27it/s, 191 steps of size 3.34e-02. acc. prob=0.94]sample:  76%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ  | 761/1000 [00:10<00:01, 134.33it/s, 63 steps of size 3.34e-02. acc. prob=0.94] sample:  78%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š  | 775/1000 [00:10<00:02, 107.48it/s, 63 steps of size 3.34e-02. acc. prob=0.94]sample:  79%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š  | 787/1000 [00:10<00:02, 94.18it/s, 63 steps of size 3.34e-02. acc. prob=0.94] sample:  80%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰  | 798/1000 [00:11<00:02, 96.16it/s, 31 steps of size 3.34e-02. acc. prob=0.94]sample:  81%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ  | 809/1000 [00:11<00:03, 60.48it/s, 63 steps of size 3.34e-02. acc. prob=0.94]sample:  82%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 818/1000 [00:11<00:03, 53.76it/s, 127 steps of size 3.34e-02. acc. prob=0.93]sample:  82%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 825/1000 [00:11<00:03, 55.36it/s, 63 steps of size 3.34e-02. acc. prob=0.94] sample:  83%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž | 834/1000 [00:11<00:02, 60.52it/s, 191 steps of size 3.34e-02. acc. prob=0.94]sample:  84%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ– | 843/1000 [00:11<00:02, 66.38it/s, 63 steps of size 3.34e-02. acc. prob=0.94] sample:  85%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 853/1000 [00:12<00:02, 72.02it/s, 63 steps of size 3.34e-02. acc. prob=0.94]sample:  86%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ | 862/1000 [00:12<00:01, 75.06it/s, 127 steps of size 3.34e-02. acc. prob=0.94]sample:  87%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹ | 871/1000 [00:12<00:01, 77.84it/s, 127 steps of size 3.34e-02. acc. prob=0.94]sample:  88%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š | 880/1000 [00:12<00:01, 79.87it/s, 63 steps of size 3.34e-02. acc. prob=0.94] sample:  89%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 889/1000 [00:12<00:01, 72.29it/s, 63 steps of size 3.34e-02. acc. prob=0.94]sample:  90%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰ | 897/1000 [00:12<00:01, 67.71it/s, 63 steps of size 3.34e-02. acc. prob=0.94]sample:  91%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ | 908/1000 [00:12<00:01, 76.35it/s, 63 steps of size 3.34e-02. acc. prob=0.94]sample:  92%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–| 918/1000 [00:12<00:00, 82.14it/s, 63 steps of size 3.34e-02. acc. prob=0.94]sample:  93%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž| 927/1000 [00:13<00:00, 74.91it/s, 191 steps of size 3.34e-02. acc. prob=0.94]sample:  94%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Ž| 935/1000 [00:13<00:01, 56.42it/s, 63 steps of size 3.34e-02. acc. prob=0.94] sample:  94%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–| 942/1000 [00:13<00:01, 50.32it/s, 63 steps of size 3.34e-02. acc. prob=0.94]sample:  95%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–| 948/1000 [00:13<00:01, 40.38it/s, 95 steps of size 3.34e-02. acc. prob=0.94]sample:  95%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ| 953/1000 [00:13<00:01, 38.38it/s, 63 steps of size 3.34e-02. acc. prob=0.94]sample:  96%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Œ| 958/1000 [00:14<00:01, 38.80it/s, 191 steps of size 3.34e-02. acc. prob=0.94]sample:  96%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹| 964/1000 [00:14<00:00, 41.99it/s, 127 steps of size 3.34e-02. acc. prob=0.94]sample:  97%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‹| 970/1000 [00:14<00:00, 45.11it/s, 127 steps of size 3.34e-02. acc. prob=0.94]sample:  98%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š| 975/1000 [00:14<00:00, 43.32it/s, 63 steps of size 3.34e-02. acc. prob=0.94] sample:  98%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š| 980/1000 [00:14<00:00, 43.30it/s, 63 steps of size 3.34e-02. acc. prob=0.94]sample:  98%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–Š| 985/1000 [00:14<00:00, 44.26it/s, 127 steps of size 3.34e-02. acc. prob=0.94]sample:  99%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰| 990/1000 [00:14<00:00, 44.15it/s, 127 steps of size 3.34e-02. acc. prob=0.94]sample: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–‰| 995/1000 [00:14<00:00, 36.85it/s, 255 steps of size 3.34e-02. acc. prob=0.94]sample: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 1000/1000 [00:14<00:00, 66.81it/s, 63 steps of size 3.34e-02. acc. prob=0.94]
arviz - WARNING - Shape validation failed: input_shape: (1, 500), minimum_shape: (chains=2, draws=4)
mean sd hdi_5.5% hdi_94.5% mcse_mean mcse_sd ess_bulk ess_tail r_hat
a 1.38 1.02 0.11 2.69 0.07 0.06 192.96 201.62 NaN
b 0.28 0.09 0.15 0.42 0.01 0.01 177.27 192.47 NaN
etasq 0.19 0.19 0.01 0.39 0.02 0.02 129.19 254.18 NaN
g 0.60 0.55 0.03 1.27 0.04 0.04 202.92 271.68 NaN
kernel[0] -0.16 0.32 -0.72 0.27 0.03 0.03 108.09 103.43 NaN
kernel[1] -0.02 0.31 -0.50 0.53 0.03 0.03 108.18 83.40 NaN
kernel[2] -0.08 0.29 -0.52 0.39 0.03 0.03 104.52 82.13 NaN
kernel[3] 0.36 0.26 -0.06 0.77 0.03 0.03 109.68 101.92 NaN
kernel[4] 0.07 0.25 -0.31 0.48 0.03 0.02 97.98 100.86 NaN
kernel[5] -0.40 0.28 -0.89 -0.00 0.03 0.02 105.36 84.04 NaN
kernel[6] 0.14 0.24 -0.23 0.56 0.02 0.02 98.67 87.33 NaN
kernel[7] -0.22 0.25 -0.58 0.19 0.03 0.02 84.80 101.42 NaN
kernel[8] 0.25 0.23 -0.14 0.59 0.03 0.02 80.88 114.93 NaN
kernel[9] -0.19 0.31 -0.63 0.34 0.03 0.01 139.57 199.16 NaN
rhosq 1.29 1.61 0.03 3.31 0.10 0.10 155.41 262.27 NaN 
library(BayesianInference)
jnp = reticulate::import('jax.numpy')
pd = reticulate::import('pandas')
# setup platform------------------------------------------------
m=importBI(platform='cpu')

# Import data ------------------------------------------------
m$data(m$load$kline2(only_path=T), sep=';')
islandsDistMatrix = m$load$islands_dist_matrix(frame = FALSE)$data
m$data_to_model(list('total_tools', 'population'))
m$data_on_model$society = jnp$arange(0,10, dtype='int64')
m$data_on_model$Dmat = jnp$array(islandsDistMatrix)


# Define model ------------------------------------------------
model <- function(Dmat, population, society, total_tools){
  a = bi.dist.exponential(1, name = 'a')
  b = bi.dist.exponential(1, name = 'b')
  g = bi.dist.exponential(1, name = 'g')
  
  # non-centered Gaussian Process prior
  etasq = bi.dist.exponential(2, name = 'etasq')
  rhosq = bi.dist.exponential(0.5, name = 'rhosq')
  k = m$gaussian$gaussian_process(Dmat, etasq, rhosq, 0.01)
  
  lambda_ = a * population**b / g * jnp$exp(k[society])
  m$dist$poisson(lambda_, obs=total_tools)
}

# Run MCMC ------------------------------------------------
m$fit(model) # Optimize model parameters through MCMC sampling

# Summary ------------------------------------------------
m$summary() # Get posterior distribution
using BayesianInference

# Setup device------------------------------------------------
m = importBI(platform="cpu")

# Import Data & Data Manipulation ------------------------------------------------
# Import
data_path = m.load.kline2(only_path = true)
m.data(data_path, sep=";") 

islandsDistMatrix = m.load.islands_dist_matrix(frame = false)["data"]
m.data_to_model(["total_tools", "population"])
m.data_on_model["society"] = jnp.arange(0,10)# index observations
m.data_on_model["Dmat"] = jnp.array(islandsDistMatrix) # Distance matrix



# Define model ------------------------------------------------
@BI function model(Dmat, population, society, total_tools)
    a = m.dist.exponential(1, name = "a")
    b = m.dist.exponential(1, name = "b")
    g = m.dist.exponential(1, name = "g")

    # non-centered Gaussian Process prior
    etasq = m.dist.exponential(2, name = "etasq")
    rhosq = m.dist.exponential(0.5, name = "rhosq")
    SIGMA = etasq * jnp.exp(-rhosq * jnp.square(Dmat))
    SIGMA = SIGMA.at[jnp.diag_indices(Dmat.shape[0])].add(etasq)
    k = m.dist.multivariate_normal(0, SIGMA, name = "k")

    lambda_ = a * population^b / g * jnp.exp(k[society])

    m.dist.poisson(lambda_, obs=total_tools)

end

# Run mcmc ------------------------------------------------
m.fit(model)  # Optimize model parameters through MCMC sampling

# Summary ------------------------------------------------
m.summary() # Get posterior distributions

Mathematical Details

Formula

The following equation allows us to evaluate the relationship between the dependent variable Y distributed normal, and the independent variable X while incorporating a GP for the effect of variable Q:

Y_{[i]} \sim \text{Normal}( \alpha + \beta X_{[i]} + \gamma_{[Q_{[i]}]}, \sigma)

where:

  • Y_{[i]} is the i-th value for the dependent variable Y.

  • \alpha is the intercept term.

  • \beta is the regression coefficient term.

  • X_{[i]} is the i-th value for the independent variable X.

  • Q_{[i]} is an integer-valued independent variable (e.g., year-of-birth, age, year) for observation i.

  • \gamma is a vector output from a Gaussian process:

\gamma \sim \text{MVNormal} \left( Z, \varsigma\Omega\varsigma \right)

where:

  • Z represents the mean vector of the multivariate normal distribution and set to zero πŸ›ˆ.

  • \varsigma is a diagonal matrix of standard deviations.

  • \Omega is a correlation matrix.

  • Multiple kernel functions for \Omega exist and will be discussed in the Note(s) section. But the most common one is the quadratic kernel:

\Omega_{[i,j]} = \eta \exp(-\phi^2 D_{[i,j]}^2)

Where:

  • \eta is the maximal correlation.

  • \phi determines the rate of decline.

  • D_{[i,j]} is the distance between the i-th and j-th categories.

Bayesian model

In the Bayesian formulation, we define each parameter with priors πŸ›ˆ. We can express a Bayesian version of this GP using the following model:

Y_i = \alpha + \beta X_i + \gamma_{Z_i}

\gamma \sim \text{MVNormal} \left( \begin{pmatrix} 0 \\ \vdots \\ 0 \end{pmatrix}, K \right)

K_{ij} = \eta^2 \exp(-p^2D_{ij}^2) + \delta_{ij} \sigma^2

\alpha \sim \text{Normal}(0,1)

\eta^2 \sim \text{HalfCauchy}(0,1)

p^2 \sim \text{HalfCauchy}(0,1)

where:

  • Y_i is the i-th value for the dependent variable Y.

  • \alpha is the intercept term with a prior of \text{Normal}(0,1).

  • \beta is the regression coefficient term with a prior of \text{Normal}(0,1).

  • X_i is the i-th value for the independent variable X.

  • \gamma_{Z_i} is the Gaussian process i-th value for the independent variable Z.

  • \gamma is the latent function modeled by the GP.

  • K_{ij} is the kernel function evaluated at the corresponding points, K_{ij} = k(Z_i, Z_j), with priors of HalfCauchy(0,1) for \eta^2 and p^2 to ensure positive values.

Notes

Note

Common kernel functions include:

  • Radial Basis Function (RBF) or Squared Exponential Kernel: k(x,x') = \sigma^2 \exp\left(-\frac{||x-x'||^2}{2l^2}\right)

  • Rational Quadratic Kernel, this kernel is equivalent to adding together many RBF kernels with different length scales: k(x,x') = \sigma^2 \left(1 + \frac{||x-x'||^2}{2l^2}\right)^{-\alpha}

  • Periodic kernel allows for modeling functions that repeat themselves exactly: k(x,x') = \sigma^2 \exp\left(-\frac{2\sin^2(\pi||x-x'||/p)}{l^2}\right)

  • Locally Periodic Kernel:

k(x,x') = \sigma^2 \exp\left(-\frac{2\sin^2(\pi||x-x'||/p)}{l^2}\right) \exp\left(-\frac{||x-x'||^2}{2l^2}\right)

  • Any slope or intercept in your model can be defined using a Gaussian Process.

Reference(s)

McElreath, Richard. 2018. Statistical Rethinking: A Bayesian course with examples in R and Stan. Chapman; Hall/CRC.

https://www.cs.toronto.edu/~duvenaud/cookbook/